from demo_qwen3_api import get_qwen3_client, get_ans_dict_by_search, get_non_empty_keys,\
    get_matched_info_by_jsondict

qwen_client = None
df=None
empty_emb=None
template_json_str=None

def search_and_show(search_txt, threshold=0.5):
    """
    检查功能说明：
    - 按照描述，从CSV读取服装数据
    - 调用Qwen3进行推理，获取JSON结构，并解析非空关键 key
    - 对key链做限定筛选，仅保留感兴趣的字段
    - 提取特征
    - 计算embedding（CLIP-ViT-B32），并与数据库内容做相似度对比
    - 多条件（刚性、embedding）排队，最终返回TopK相似图片及内容
    """

    import pandas as pd
    import json
    import heapq
    import torch
    import torch.nn.functional as F
    import numpy as np
    from copy import deepcopy as dcp
    from util_for_clipvitb32 import clip_vit_b32, get_embedding
    from utils.util_flux import process_img_1024

    global qwen_client, df, empty_emb, template_json_str

    # 读取csv数据，只读一次
    if df is None:
        # csv_file_path = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data/qwenvl3_cloth_data2600.csv'
        csv_file_path = '/mnt/nas/shengjie/qdrant_data/qwenvl3_data/qwenvl3_cloth_data_api.csv'
        df = pd.read_csv(csv_file_path)

    # 读取模板json字符串，只读一次
    if template_json_str is None:
        # 读取模板json字符串
        template_json_path = 'demo_qwenvl_clothing.json'
        with open(template_json_path, 'r', encoding='utf-8') as f:
            template_json_str = f.read()

    # 获取qwen3客户端，只获取一次
    if qwen_client is None:
        qwen_client = get_qwen3_client()

    # 获取推理结果
    ans = get_ans_dict_by_search(qwen_client, template_json_str, search_txt)

    try:
        ans_dict = json.loads(ans)
    except Exception as e:
        print("Failed to parse ans as JSON:", e)
        ans_dict = None

    # 获取所有非空key链
    non_empty_keys = get_non_empty_keys(ans_dict)
    '''
    key 需要做筛选
    1 只筛选 
        basic_info : season, gender     e.g. ['basic_info','season'] ['basic_info','season']
        design: color.primary color.secondary color.pattern
                category.main category.sub
                sleeve.type sleeve.cuff sleeve.count e.g. ['design', 'sleeve', 'type']...
                collar.type
                pockets
        material: composition
        decoration
        visual
        # 最后的keys只包含以上的key序列
        # 当前 non_empty_keys 中 包含 [ ['design', 'color', 'primary'],[],...  ]
    2 数量 软对比 sim=target_num/num if num >= target_num and num != -1 else 0
    3 正反面 刚性对比
            visual.view_types 中 sim=1 search_type in types else 0
    4 search value 时 target value为空时 加几分？
    '''
    # INSERT_YOUR_CODE

    # 允许保留的key链前缀
    allowed_prefixes = [
        ['basic_info', 'season'],
        ['basic_info', 'gender'],
        ['design', 'color', 'primary'],
        ['design', 'color', 'secondary'],
        ['design', 'color', 'pattern'],
        ['design', 'category', 'main'],
        ['design', 'category', 'sub'],
        ['design', 'sleeve', 'type'],
        ['design', 'sleeve', 'cuff'],
        ['design', 'sleeve', 'count'],
        ['design', 'collar', 'type'],
        ['design', 'pockets'],
        # ['material', 'composition'],
        ['decoration'],
        ['visual'],
    ]

    # 给定实际key链，判定是否属于允许集合
    def is_key_allowed(keys):
        # 判定当前keys链是否在可允许的prefix集合内
        for prefix in allowed_prefixes:
            if len(keys) < len(prefix):
                continue
            # 只要前缀对上即可
            if keys[:len(prefix)] == prefix:
                return True
        return False

    print('[before] non_empty_keys:', non_empty_keys)
    # 保留非空且受允许key链的字段
    non_empty_keys = [keys for keys in non_empty_keys if is_key_allowed(keys)]
    print('[sifted] non_empty_keys:', non_empty_keys)

    # 获取embedding模型
    clip_model = clip_vit_b32()
    if empty_emb is None:
        empty_txt = ''
        empty_emb = get_embedding(clip_model, empty_txt)

    # 获得本次搜索条件的对应值及嵌入向量
    search_joined_values_list = get_matched_info_by_jsondict(ans_dict, non_empty_keys)
    search_joined_values_embs = get_embedding(clip_model, search_joined_values_list) # K' 512
    search_joined_values = '&&'.join(search_joined_values_list)                                    

    topk = 10
    heap = []

    '''
    1 初筛 刚性条件
    2 sim筛 凑齐所有的value(K*N个) 一起计算emb (K*N,512) => (N,K,512)
    3 search_emb (1,K,512).expand->(N,K,512) , cossim( value_emb, search_emb ) => sim_matrix (N,K)
    4 sim_matrix 则 [ [k1,k2,k3..]...[] ] 筛除 ki < threshold的数据 => sifted_sim_matrix (N' , K)
    5 对每个 sum(N) 排序 取 topk
    '''

    sifted_data = []
    sifted_value = []
    sifted_search_emb = []
    sifted_non_empty_keys = dcp(non_empty_keys)

    # 1 初筛 刚性条件
    for idx, row in df.iterrows():
        # target_json_dict_str = row.get('desc_json', None)  # for local
        target_json_dict_str = row.get('desc_dict', None)  # for api
        img_path = row.get('img_path', '')

        try:
            import ast
            target_json_dict = ast.literal_eval(target_json_dict_str)
            # target_json_dict = json.loads(target_json_dict_str)
        except Exception as e:
            # print("template_json_str 解析失败:", e)
            target_json_dict = {}

        joined_values_list = get_matched_info_by_jsondict(target_json_dict, non_empty_keys)
        
        # 对 每一个 item 都做 相似度，每个相似度 when if < thre，则continue
        # 否则 相似度累加 sum_sim , 按照这个排序
        sim_dict = {}
        satisfy = True
        for joined_value, search_joined_value,cur_keys in \
            zip(joined_values_list, search_joined_values_list,non_empty_keys):
            key_str = '->'.join(map(str,cur_keys))

            if search_joined_value.isdigit() or \
                (search_joined_value.startswith('-') and search_joined_value[1:].isdigit()):
                search_num = int(search_joined_value)
                # 检查数值刚性筛选
                if search_num == 0:
                    # 比如口袋数为0（要求无口袋），target无返回空str；等于""
                    if joined_value == "":
                        sim_dict[key_str] = 1
                        continue
                    # 检索 要求 没有, 则 target 也得 没有,此时target 必须为 -1
                    # else satisfy=False
                    if joined_value.isdigit() or (joined_value.startswith('-') and joined_value[1:].isdigit()):
                        joined_num = int(joined_value)
                        # 检索要求没有，目标只能是-1或0
                        if joined_num in [0, -1]:
                            sim_dict[key_str] = 1
                            continue
                        else:
                            satisfy = False
                            break
                    else:
                        satisfy = False
                        break
                # search num != 0
                if joined_value.isdigit() or (joined_value.startswith('-') and joined_value[1:].isdigit()):
                    joined_num = int(joined_value)

                    # sim = search_num / joined_num if joined_num >= search_num and joined_num != -1 else 0
                    sim = 1 if search_num == joined_num else 0
                    # if sim == 0:
                    if sim < threshold:
                        satisfy = False
                        break
                    sim_dict[key_str] = sim
                    continue
                else:
                    satisfy = False
                    break

            # cur_keys = non_empty_keys[joined_values_list.index(joined_value)] if joined_value in joined_values_list else []
            # 判断key链是否包含visual和view_types
            if len(cur_keys) >= 2 and cur_keys[-2:] == ["visual", "view_types"]:
                sv_set = set([v.strip() for v in str(search_joined_value).split(",")])
                jv_set = set([v.strip() for v in str(joined_value).split(",")])
                # 检索的视角数量可能大于 json 的视角数量
                # if not sv_set.issubset(jv_set):
                if not jv_set.issubset(sv_set):
                    satisfy = False
                    break
                # 逻辑过了，给sim=1
                sim_dict[key_str] = 1
                continue

        # 满足刚性筛选的数据进入后续步骤
        if satisfy:
            sifted_data.append(
                (
                    img_path,
                    joined_values_list,
                    sim_dict,
                )
            )
            # 对剩下的维度做多模态相似度
            for joined_value, search_emb, cur_keys, search_joined_value in \
                zip(joined_values_list, search_joined_values_embs, non_empty_keys, search_joined_values_list):
                # 不对刚性或数值类、视角类的做sim embedding计算
                skip_numeric = (
                    (
                        search_joined_value.isdigit() or 
                        (search_joined_value.startswith('-') and search_joined_value[1:].isdigit())
                    )
                )
                skip_visual = (len(cur_keys) >= 2 and cur_keys[-2:] == ["visual", "view_types"])
                if skip_numeric or skip_visual:
                    if cur_keys in sifted_non_empty_keys:
                        sifted_non_empty_keys.remove(cur_keys)
                    continue
                sifted_value.append(joined_value)
                sifted_search_emb.append(search_emb)

    # ========= 多模态embedding筛选 =========
    if len(sifted_value):
        # 2 sim筛 凑齐所有的value(K*N个) 一起计算emb (K*N,512) => (N,K,512)
        value_emb = get_embedding(clip_model, sifted_value) # numpy N*K, 512
        N = len(sifted_data)
        value_emb = value_emb.reshape(N,-1,512) # N K 512

        # expand search_emb
        # 3 search_emb (1,K,512).expand->(N,K,512) , cossim( value_emb, search_emb ) => sim_matrix (N,K)
        # sifted_search_emb list N*K, 512
        search_emb_expanded = np.array(sifted_search_emb)  # numpy shape: (N*K, 512)
        search_emb = search_emb_expanded.reshape(N,-1,512)
        search_sim_matrix = F.cosine_similarity(
                    torch.tensor(value_emb), # torch.Size([1117, 4, 512])
                    torch.tensor(search_emb), # (1117, 4, 512)
                    dim=-1
                ) # N,K  torch.Size([1117, 4])
        
        
        # 4 sim_matrix 则 [ [k1,k2,k3..]...[] ] 筛除 ki < threshold的数据 => sifted_sim_matrix (N' , K)
        # 返回筛除数据的下标
        # 构建threshold矩阵, 矩阵运算一次获取结果,结果中呈现负值的就是被筛除的,获得哪些数据所在的行
        # threshold = 0.4  # 按需可设置
        # search_sim_matrix: (N, K) tensor 或 ndarray

        # 将其变成numpy确保可以矩阵运算
        if hasattr(search_sim_matrix, "detach"):  # torch tensor
            sim_mat_np = search_sim_matrix.detach().cpu().numpy()
        else:
            sim_mat_np = np.asarray(search_sim_matrix)
        
        # threshold矩阵，与sim_mat_np一样shape
        threshold_mat = np.full_like(sim_mat_np, threshold)
        judged = sim_mat_np - threshold_mat  # 大于0表示通过，小于0表示未通过
        valid_row_mask = np.all(judged >= 0, axis=1)  # shape: (N,), True为整行都通过
        keep_indices = np.where(valid_row_mask)[0].tolist()
        # sifted_sim_matrix = sim_mat_np[valid_row_mask]
        total_sift_num = len(keep_indices)

        # keep_indices是保留行下标，sifted_sim_matrix为对应的sim子矩阵


        # 5 对每个 sum(N) 排序 取 topk
        # sifted_non_empty_keys = non_empty_keys - 
        for is_satisfy, search_sim, data in zip(valid_row_mask,search_sim_matrix, sifted_data):
            if not is_satisfy : continue
            img_path, joined_values_list, sim_dict = data
            # search_sim是K维tensor，赋分sifted_non_empty_keys
            # 注意cur_keys赋值在循环外有坑，用临时变量
            for cur_key, sim_score in zip(sifted_non_empty_keys, search_sim.tolist()):
                key_str = '->'.join(cur_key)
                sim_dict[key_str] = sim_score

            # sim_list = list(sim_dict.values())
            sim_list = []
            for k in non_empty_keys:
                key_str = '->'.join(map(str,k))
                sim_list.append( sim_dict[key_str] )
            sum_sim = sum(sim_list)
            item = (sum_sim, [f'{v:.4f}' for v in sim_list], img_path, '&&'.join(joined_values_list))
            if len(heap) < topk:
                heapq.heappush(heap, item)
            else:
                if sum_sim > heap[0][0]:
                    heapq.heappushpop(heap, item)
    else:
        # 若无多模态维度，直接对刚性过筛项目排序/取TopK
        total_sift_num = len(sifted_data)
        for data in sifted_data:
            img_path, joined_values_list, sim_dict = data
            sim_list = list(sim_dict.values())
            sum_sim = sum(sim_list)
            item = (sum_sim, [f'{v:.4f}' for v in sim_list], img_path, '&&'.join(joined_values_list))
            if len(heap) < topk:
                heapq.heappush(heap, item)
            else:
                if sum_sim > heap[0][0]:
                    heapq.heappushpop(heap, item)

    print('筛选后的总数据量: ',total_sift_num)
    # ================ 排序+准备输出 =====================
    heap_sorted = sorted(heap, key=lambda x: -x[0])
    res_list = []
    img_list = []
    for sim, sim_list, path, tar_value_str in heap_sorted:
        res_list.append(f"TOPK sum_sim={sim:.4f} search_sim={sim_list}  path= {path} desc={tar_value_str}")
        img_list.append(process_img_1024(path))

    # 格式化search内容列表
    search_txt_list = [
        f"{'->'.join(map(str, k))}: {v}"
        for k, v in zip(non_empty_keys, search_joined_values_list)
    ]
    search_total_txt = f"筛选后的总数据量: {total_sift_num}\n"
    search_txt_show = search_total_txt + '\n'.join(search_txt_list)
    return search_txt_show, img_list, '\n'.join(res_list)

def start_gradio(port):
    import gradio as gr
    with gr.Blocks() as demo:
        gr.Markdown("# QwenVL3 Gradio Demo (以文搜图)")
        with gr.Row():
            # with gr.Column():
            search_box = gr.Textbox(label="输入搜索描述", placeholder="如：春秋两季黑色大衣，两个口袋", lines=2)
        with gr.Row():
            threshold_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.01, label="阈值 threshold", interactive=True)
            submit_btn = gr.Button("开始搜索")
            # with gr.Column():
        with gr.Row():
            output_search = gr.Textbox(label="搜索内容", interactive=False)
        with gr.Row():
            gallery = gr.Gallery(label="TopK图片预览", columns=5, height="auto")
        with gr.Row():
            output_list = gr.Textbox(label="TopK内容", lines=10, interactive=False)
        submit_btn.click(
            search_and_show,
            inputs=[search_box, threshold_slider],
            outputs=[output_search, gallery, output_list]
        )
    demo.launch(server_name="0.0.0.0", server_port=port)

# 若本py被直接调用则自动启动界面
if __name__ == '__main__':
    import os
    os.environ['GRADIO_TEMP_DIR'] = '/mnt/nas/shengjie/tmp'
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
    parser.add_argument('-p', '--port', type=int, default=20023, help='CUDA device id')
    args, unknown = parser.parse_known_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

    start_gradio(args.port)